﻿using ServiceStack.Logging;
using System;
using System.Collections.Generic;
using System.IO;
using System.Security.Cryptography.X509Certificates;
using System.Security.Cryptography.Xml;
using System.Xml;

namespace OneLoginPlugin
{
    public class Certificate
    {
        public X509Certificate2 cert;

        public void LoadCertificate(string certificate)
        {
            cert = new X509Certificate2();
            cert.Import(StringToByteArray(certificate));
        }

        public void LoadCertificate(byte[] certificate)
        {
            cert = new X509Certificate2();
            cert.Import(certificate);
        }

        private byte[] StringToByteArray(string st)
        {
            byte[] bytes = new byte[st.Length];
            for (int i = 0; i < st.Length; i++)
            {
                bytes[i] = (byte)st[i];
            }
            return bytes;
        }
    }

    public class Response
    {
        private XmlDocument xmlDoc;
        private AccountSettings accountSettings;
        private Certificate certificate;

        public const string SAMLNamespaceUrl = "urn:oasis:names:tc:SAML:2.0:assertion";
        public const string SAMLPNamespaceUrl = "urn:oasis:names:tc:SAML:2.0:protocol";

        public string Xml { get; set; }

        private static readonly ILog _log = LogManager.GetLogger(typeof(Response));


        public Response(AccountSettings accountSettings)
        {
            this.accountSettings = accountSettings;
            certificate = new Certificate();
            certificate.LoadCertificate(accountSettings.certificate);
        }

        public void LoadXml(string xml)
        {
            this.Xml = xml;
            xmlDoc = new XmlDocument();
            xmlDoc.PreserveWhitespace = true;
            xmlDoc.XmlResolver = null;
            xmlDoc.LoadXml(xml);
        }

        public void LoadXmlFromBase64(string response)
        {
            System.Text.ASCIIEncoding enc = new System.Text.ASCIIEncoding();
            LoadXml(enc.GetString(Convert.FromBase64String(response)));
        }


        private XmlNodeList getNodes(string xpath)
        {
            XmlNamespaceManager manager = new XmlNamespaceManager(xmlDoc.NameTable);
            manager.AddNamespace("saml", SAMLNamespaceUrl);
            manager.AddNamespace("samlp", SAMLPNamespaceUrl);
            return xmlDoc.SelectNodes(xpath, manager);
        }

        /// <summary>
        /// Pass a list of methods to test an attribute.  The method should return true if it fails, otherwise return false;
        /// </summary>
        /// <param name="xpath"></param>
        /// <param name="attributeFailTests"></param>
        /// <returns></returns>
        private bool TestAttributesTrue(string xpath, params Func<XmlAttribute, bool>[] attributeFailTests)
        {

            foreach (XmlNode node in getNodes(xpath))
            {
                foreach (XmlAttribute attribute in node.Attributes)
                {
                    foreach (var attributeTest in attributeFailTests)
                    {
                        if (attributeTest(attribute) == true)
                        {
                            return false;
                        }
                    }
                }
            }

            return true;
        }

        public bool IsValidSubjectConfirmationData()
        {
            foreach (XmlNode node in getNodes("//saml:SubjectConfirmationData"))
            {
                foreach (XmlAttribute attribute in node.Attributes)
                {

                    if (attribute.Name == "NotOnOrAfter")
                    {
                        DateTime notOnOrafter = DateTime.Parse(attribute.Value).ToUniversalTime();
                        if (DateTime.UtcNow > notOnOrafter)
                        {
                            _log.DebugFormat("Authentication failed on NotOnOrAfter: {0}", notOnOrafter);
                            return false;
                        }
                    }

                    if (attribute.Name == "Recipient")
                    {

                        if (accountSettings.Recipient != attribute.Value)
                        {
                            _log.DebugFormat("Authentication failed on Recipient: {0} -- {1}", accountSettings.Recipient, attribute.Value);
                            return false;
                        }
                    }


                    if (attribute.Name == "InResponseTo")
                    {
                        if (!AuthRequest.IDs.Contains(attribute.Value))
                        {
                            _log.DebugFormat("Authentication failed on InResponseTo: {0}", attribute.Value);
                            return false;
                        }
                    }

                }
            }
            return true;
        }


        public bool IsValid()
        {
            bool status = false;

            if (IsValidSubjectConfirmationData())
            {
                XmlNamespaceManager manager = new XmlNamespaceManager(xmlDoc.NameTable);
                manager.AddNamespace("ds", SignedXml.XmlDsigNamespaceUrl);
                XmlNodeList nodeList = xmlDoc.SelectNodes("//ds:Signature", manager);

                SignedXml signedXml = new SignedXml(xmlDoc);
                foreach (XmlNode node in nodeList)
                {
                    signedXml.LoadXml((XmlElement)node);
                    status = signedXml.CheckSignature(certificate.cert, true);
                    if (!status)
                    {
                        return false;
                    }
                }
            }
            return status;
        }

        public string GetNameID()
        {
            XmlNamespaceManager manager = new XmlNamespaceManager(xmlDoc.NameTable);
            manager.AddNamespace("ds", SignedXml.XmlDsigNamespaceUrl);
            manager.AddNamespace("saml", "urn:oasis:names:tc:SAML:2.0:assertion");
            manager.AddNamespace("samlp", "urn:oasis:names:tc:SAML:2.0:protocol");

            XmlNode node = xmlDoc.SelectSingleNode("/samlp:Response/saml:Assertion/saml:Subject/saml:NameID", manager);
            return node.InnerText;
        }
    }

    public class AuthRequest
    {
        public string id;
        private string issue_instant;
        private AccountSettings accountSettings;

        public static HashSet<string> IDs = new HashSet<string>();

        public enum AuthRequestFormat
        {
            Base64 = 1
        }

        public AuthRequest(AccountSettings accountSettings)
        {
            this.accountSettings = accountSettings;

            id = "_" + System.Guid.NewGuid().ToString();

            IDs.Add(id);

            issue_instant = DateTime.Now.ToUniversalTime().ToString("yyyy-MM-ddTHH:mm:ssZ");
        }

        public string GetRequest(AuthRequestFormat format)
        {
            using (StringWriter sw = new StringWriter())
            {
                XmlWriterSettings xws = new XmlWriterSettings();
                xws.OmitXmlDeclaration = true;

                using (XmlWriter xw = XmlWriter.Create(sw, xws))
                {
                    xw.WriteStartElement("samlp", "AuthnRequest", "urn:oasis:names:tc:SAML:2.0:protocol");
                    xw.WriteAttributeString("ID", id);
                    xw.WriteAttributeString("Version", "2.0");
                    xw.WriteAttributeString("IssueInstant", issue_instant);
                    xw.WriteAttributeString("Destination", accountSettings.idp_sso_target_url);
                    xw.WriteAttributeString("ProtocolBinding", "urn:oasis:names:tc:SAML:2.0:bindings:HTTP-POST");
                    xw.WriteAttributeString("AssertionConsumerServiceURL", accountSettings.AssertionConsumerServiceUrl);

                    xw.WriteStartElement("saml", "Issuer", "urn:oasis:names:tc:SAML:2.0:assertion");
                    xw.WriteString(accountSettings.issuer);
                    xw.WriteEndElement();

                    xw.WriteStartElement("samlp", "NameIDPolicy", "urn:oasis:names:tc:SAML:2.0:protocol");
                    xw.WriteAttributeString("Format", "urn:oasis:names:tc:SAML:2.0:nameid-format:unspecified");
                    xw.WriteAttributeString("AllowCreate", "true");
                    xw.WriteEndElement();

                    xw.WriteStartElement("samlp", "RequestedAuthnContext", "urn:oasis:names:tc:SAML:2.0:protocol");
                    xw.WriteAttributeString("Comparison", "exact");
                    xw.WriteEndElement();

                    xw.WriteStartElement("saml", "AuthnContextClassRef", "urn:oasis:names:tc:SAML:2.0:assertion");
                    xw.WriteString("urn:oasis:names:tc:SAML:2.0:ac:classes:PasswordProtectedTransport");
                    xw.WriteEndElement();

                    xw.WriteEndElement();
                }

                if (format == AuthRequestFormat.Base64)
                {
                    string xx = sw.ToString();
                    byte[] toEncodeAsBytes = System.Text.ASCIIEncoding.ASCII.GetBytes(sw.ToString());
                    return System.Convert.ToBase64String(toEncodeAsBytes);
                }

                return null;
            }
        }

        public string GetLogoutRequest(AuthRequestFormat format, string userId)
        {
            using (StringWriter sw = new StringWriter())
            {
                XmlWriterSettings xws = new XmlWriterSettings();
                xws.OmitXmlDeclaration = true;

                using (XmlWriter xw = XmlWriter.Create(sw, xws))
                {
                    xw.WriteStartElement("samlp", "LogoutRequest", "urn:oasis:names:tc:SAML:2.0:protocol");
                    xw.WriteAttributeString("ID", id);
                    xw.WriteAttributeString("Version", "2.0");
                    xw.WriteAttributeString("IssueInstant", issue_instant);
                    xw.WriteAttributeString("Destination", accountSettings.idp_sso_target_url);

                    xw.WriteStartElement("saml", "Issuer", "urn:oasis:names:tc:SAML:2.0:assertion");
                    xw.WriteString(accountSettings.issuer);
                    xw.WriteEndElement();

                    xw.WriteStartElement("saml", "NameID", "urn:oasis:names:tc:SAML:2.0:assertion");
                    xw.WriteAttributeString("Format", "urn:oasis:names:tc:SAML:1.1:nameid-format:emailAddress");
                    xw.WriteString("david.churchland@hp.com");
                    xw.WriteEndElement();

                    xw.WriteEndElement();
                }

                if (format == AuthRequestFormat.Base64)
                {
                    string xx = sw.ToString();
                    byte[] toEncodeAsBytes = System.Text.ASCIIEncoding.ASCII.GetBytes(sw.ToString());
                    return System.Convert.ToBase64String(toEncodeAsBytes);
                }

                return null;
            }
        }
    }
}
